import torch
from torch import nn
import torchvision
from torch.utils.data import DataLoader

# 这个数据集100多G

#train_data = torchvision.datasets.ImageNet("./imgnetdata", split='train', download=True, transform=torchvision.transforms.ToTensor())

vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True)

#vgg16是一个分类模型，可以分为1000类

#cifar是一个分为10类的分类器
dataset = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
dataloader = DataLoader(dataset, batch_size= 1)

#对现有的网络模型vgg16进行改动，1000 => 10

vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10))


